'''
Author: tomwoo tom.woo@outlook.com
Date: 2025-07-24 14:22:18
LastEditors: tomwoo tom.woo@outlook.com
LastEditTime: 2025-07-25 17:28:12
FilePath: /Multi-Modal_AI_Agents/vss_agent.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''

import requests


STATUS_OK = 0
STATUS_ERR = -1


class VssAgent:
    def __init__(self, host, port):
        vss_url = f"http://{host}:{port}"

        self.health_endpoint = vss_url + "/health/ready" #check the status of the VSS server
        self.models_endpoint = vss_url + "/models" #view the configured LLM in VSS
        self.files_endpoint = vss_url + "/files" #upload and manage files
        self.summarize_endpoint = vss_url + "/summarize" #summarize uploaded content
        self.qna_endpoint = vss_url + "/chat/completions" # ask questions about uploaded content

        self.video_id = None

    def clear(self):
        self.video_id = None

    def check(self):
        try:
            response = requests.get(self.health_endpoint)

            if response.status_code == 200:
                return STATUS_OK, f"Service is ready."
            else:
                return STATUS_ERR, f"Service is not ready. Status code: {response.status_code}"

        except requests.RequestException as e:
            return STATUS_ERR, f"An error occurred: {e}"

    def upload(self, video_filename):
        if video_filename is None:
            return STATUS_ERR, "No video file provided."

        files = {'file': ('filename_with_extension', open(video_filename, 'rb'))}

        try:
            response = requests.post(self.files_endpoint, data={'purpose': 'vision', 'media_type': 'video'}, files=files)

            if response.status_code == 200:
                response_data = response.json()
                # Store the video ID
                self.video_id = response_data.get('id')
                return (STATUS_OK, self.video_id) if self.video_id else (STATUS_ERR, "No ID received.")
            else:
                return STATUS_ERR, f"Failed to upload file. Status code: {response.status_code}"

        except requests.RequestException as e:
            return STATUS_ERR, f"An error occurred: {e}"

    def summarize(self, prompt, caption_summarization_prompt, summary_aggregation_prompt):
        if self.video_id is None:
            return STATUS_ERR, "No video uploaded."

        body = {
            "id": self.video_id, # ID of file returned after upload
            "prompt": prompt,
            "caption_summarization_prompt": caption_summarization_prompt,
            "summary_aggregation_prompt": summary_aggregation_prompt,
            "model": "vila-1.5",
            "max_tokens": 1024,
            "temperature": 0.2,
            "top_p": 0.8,
            "chunk_duration": 10,
            "chunk_overlap_duration": 0,
            "enable_chat": True,
        }

        try:
            response = requests.post(self.summarize_endpoint, json=body)

            if response.status_code == 200:
                response_data = response.json()
                # Extracting the summary content
                summary = response_data.get('choices', [])[0].get('message', {}).get('content', '')
                return (STATUS_OK, summary) if summary else (STATUS_ERR, "No summary received.")
            else:
                return STATUS_ERR, f"Failed to summarize. Status code: {response.status_code}. Response: {response.text}"

        except requests.RequestException as e:
            return STATUS_ERR, f"An error occurred: {e}"

    def qna(self, question):
        if self.video_id is None:
            return STATUS_ERR, "No video uploaded."

        payload = {
            "id": self.video_id, # ID of file returned after upload
            "messages": [
                {
                    "content": question,
                    "role": "user",
                }
            ],
            "model": "vila-1.5"
        }

        try:
            response = requests.post(self.qna_endpoint, json=payload)

            if response.status_code == 200:
                response_data = response.json()
                # Extracting the answer content
                answer = response_data.get("choices", [])[0].get("message", {}).get("content", "")
                return (STATUS_OK, answer) if answer else (STATUS_ERR, "No answer received.")
            else:
                return STATUS_ERR, f"Failed to get an answer. Status code: {response.status_code}. Response: {response.text}"

        except requests.RequestException as e:
            return STATUS_ERR, f"An error occurred: {e}"

# end of file
